import torchvision.models as models
import torch.nn as nn
import torch
from torchvision.models import vgg16

class ResNet18(nn.Module):
    def __init__(self, num_classes):
        super(ResNet18, self).__init__()
        # 使用 torchvision 中的 ResNet-18 预训练模型
        self.resnet18 = models.resnet18(pretrained=True)
        # 替换最后一层全连接层，输出为 num_classes 类别
        self.resnet18.fc = nn.Linear(self.resnet18.fc.in_features, num_classes)

    def forward(self, x):
        return self.resnet18(x)

class MobileNetV2Model(nn.Module):
    def __init__(self, num_classes):
        super(MobileNetV2Model, self).__init__()
        self.model = models.mobilenet_v2(pretrained=True)
        # 修改最后一层输出类别数
        self.model.classifier[1] = nn.Linear(self.model.classifier[1].in_features, num_classes)

    def forward(self, x):
        return self.model(x)

class SqueezeNetModel(nn.Module):
    def __init__(self, num_classes):
        super(SqueezeNetModel, self).__init__()
        self.model = models.squeezenet1_1(pretrained=True)
        # 修改最后一层输出类别数
        self.model.classifier[1] = nn.Conv2d(512, num_classes, kernel_size=(1, 1))

    def forward(self, x):
        return self.model(x)

class DenseNet69(nn.Module): #这个我试过了，不管怎么调size都不行
    def __init__(self, num_classes, pretrained=False):
        super(DenseNet69, self).__init__()
        # 根据 pretrained 参数来选择是否加载预训练权重
        self.densenet = models.densenet121(pretrained=pretrained)  # 这里使用 densenet121，因为 torchvision 没有 densenet69
        self.densenet.classifier = nn.Linear(self.densenet.classifier.in_features, num_classes)

    def forward(self, x):
        return self.densenet(x)

class VGG16Model(nn.Module):
    def __init__(self, num_classes=25):
        super(VGG16Model, self).__init__()
        # 加载预训练的 VGG16 主干
        self.backbone = vgg16(pretrained=True)
        # 替换最后的分类层
        self.backbone.classifier[6] = nn.Linear(4096, num_classes)

    def forward(self, x):
        return self.backbone(x)